package com.opengate.base.servlet;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;

/**
 * @author 曳戈泰尔
 * @version 1.0
 * @description 线程安全，每个请求线程独立
 *     <p>BufferedServletRequestWrapper 是一个 ServletRequest 的包装类，用于将 ServletRequest
 *     中的数据缓存起来，以便后续多次读取数据。它的主要作用是解决在 Servlet 中多次读取请求体时，会出现 IO 异常的问题。
 *     <p>使用BufferedServletRequestWrapper可以提高性能，减少重复读取请求流的问题，适用于需要多次读取请求体的情况。
 * @date 2023/4/26 5:46 PM
 */
@Slf4j
public class BufferedServletRequestWrapper extends HttpServletRequestWrapper {

    private static final String FORWARDED = "Forwarded";
    private final ByteArrayOutputStream cacheContent;
    private final String body;

    public String getBody() {
        return body;
    }

    public BufferedServletRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        int length = Math.max(request.getContentLength(), 0);
        ByteArrayOutputStream content = new ByteArrayOutputStream(length);
        ServletInputStream is = request.getInputStream();
        byte[] buff = new byte[1024];
        int read;
        while ((read = is.read(buff)) > -1) {
            content.write(buff, 0, read);
        }
        content.flush();
        cacheContent = content;
        content.close();
        String encoding = "UTF-8";
        body = cacheContent.toString(encoding);
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        return new BufferedServletInputStream(this.getRequest().getInputStream());
    }

    @Override
    public String getHeader(String name) {
        // 先从原本的Request中获取头，如果为空且名字为token，则从参数中查找并返回
        String superHeader = super.getHeader(name);
        if (FORWARDED.equalsIgnoreCase(name) && StringUtils.isNoneBlank(superHeader)) {
            superHeader = superHeader.replaceAll("\"", "");
            return superHeader;
        }
        return superHeader;
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    private class BufferedServletInputStream extends ServletInputStream {

        private final ServletInputStream is;
        private final ByteArrayInputStream bais;

        public BufferedServletInputStream(ServletInputStream is) {
            this.is = is;
            bais = new ByteArrayInputStream(cacheContent.toByteArray());
        }

        @Override
        public boolean isFinished() {
            return this.is.isFinished();
        }

        @Override
        public boolean isReady() {
            return this.is.isReady();
        }

        @Override
        public void setReadListener(ReadListener readListener) {
            this.is.setReadListener(readListener);
        }

        @Override
        public int read() {
            return bais.read();
        }
    }
}
